package algorithm;

import dataset.Dataset;
import exception.KNoMeaningException;
import scicalc.Calculate;
import scicalc.Matrix;

public class GaussianMixture extends AbstractClusterAlgorithm {

    private int maxIter = 13;

    private int k = 3;

    private double[] alpha;

    private double[][] mu;

    private double[][][] cov;

    private double[][] gamma;

    @Override
    public long fit(Dataset X) {
        long start = System.currentTimeMillis();
        this.alpha = new double[this.k];
        for (int i = 0; i < this.k; i++)
            this.alpha[i] = Calculate.ONE / this.k;
        int[] r = Calculate.randomIntArray(0, X.shape[0], this.k);
        this.mu = Calculate.choose(X.dataset, r);
        this.cov = new double[this.k][X.shape[1]][X.shape[1]];
        for (int i = 0; i < this.k; i++)
            this.cov[i] = Calculate.multi(Matrix.identityMatrix(X.shape[1]), 0.1);
        this.gamma = new double[X.shape[0]][this.k];

        for (int iter = 0; iter < maxIter; iter++) {
            for (int j = 0; j < X.shape[0]; j++) {
                double[] alpha_p = new double[this.k];
                double sum = 0;
                for (int i = 0; i < this.k; i++) {
                    alpha_p[i] = this.alpha[i] * probability(X.dataset[j], this.mu[i], this.cov[i]);
                    sum += alpha_p[i];
                }
                this.gamma[j] = Calculate.divide(alpha_p, sum);
            }
            double[] sum_gamma_i = Calculate.sum(this.gamma, Calculate.AXIS_0);
            double[][] mu_mid = Matrix.T(Matrix.dot(Matrix.T(X.dataset), gamma));
            this.alpha = Calculate.divide(sum_gamma_i, X.shape[0]);
            for (int i = 0; i < this.k; i++) {
                this.mu[i] = Calculate.divide(mu_mid[i], sum_gamma_i[i]);
                double[][] numerator = new double[X.shape[1]][X.shape[1]];
                for (int j = 0; j < X.shape[0]; j++) {
                    double[][] xmu = Matrix.arrayToMatrix(Calculate.minus(X.dataset[j], this.mu[i]));
                    numerator = Calculate.add(numerator, Calculate.multi(Matrix.dot(Matrix.T(xmu), xmu), this.gamma[j][i]));
                }
                this.cov[i] = Calculate.divide(numerator, sum_gamma_i[i]);
            }
        }

        this.labels = new int[X.shape[0]];
        for (int j = 0; j < X.shape[0]; j++) {
            double[] alpha_p = new double[this.k];
            double sum = 0;
            for (int i = 0; i < this.k; i++) {
                alpha_p[i] = this.alpha[i] * probability(X.dataset[j], this.mu[i], this.cov[i]);
                sum += alpha_p[i];
            }
            this.labels[j] = Calculate.argmax(Calculate.divide(alpha_p, sum));
        }
        this.result.shape[0] = X.shape[0];
        this.result.shape[1] = X.shape[1];
        this.result.dataset = X.dataset;
        this.result.label = this.labels;
        return System.currentTimeMillis() - start;
    }

    private double probability(double[] x, double[] mu, double[][] cov) {
        double det = Matrix.det(cov);
        double[][] inv = Matrix.inv(cov);
        double[][] x_mu = Matrix.T(Matrix.arrayToMatrix(Calculate.minus(x, mu)));
        double exp = -0.5 * Matrix.dot(Matrix.dot(Matrix.T(x_mu), inv), x_mu)[0][0];
        double p = Calculate.exp(exp) / (Calculate.pow(2 * Calculate.PI, x.length / 2) * Calculate.sqrt(det));
        return p;
    }

    public void setK(int k) {
        if (k < 2) {
            try {
                throw new KNoMeaningException(k);
            } catch (KNoMeaningException e) {
                e.printStackTrace();
            }
        } else this.k = k;
    }

    public void setMaxIter(int maxIter) {
        this.maxIter = maxIter;
    }
}
